use proc_macro::TokenStream;

use quote::quote;
use syn::{self, AttributeArgs, Data, DeriveInput, Field, Fields, parse_macro_input};
use syn::punctuated::Punctuated;

/// 自定义派生属性
/// 测试代码表示，打算实现一个 自定义派生属性#[derive(A)]，
/// 然后为单元结构体 A 自动实 现一个实例方法 a。在调用方法 a 的时候，输出指定的字符串
#[proc_macro_derive(A)]
pub fn derive_a(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    format!(
        r#"
        impl {} {{
            pub fn test(&self) {{
                println!("test");
            }}
        }}
    "#,
        input.ident.to_string()
    )
        .parse()
        .unwrap()
}

/// 自定义属性
#[proc_macro_attribute]
pub fn to_say(args: TokenStream, input: TokenStream) -> TokenStream {
    let args = args.to_string();
    let input_str = input.to_string();
    println!("args = {:?} input = {}", args, input_str);
    format!(
        r#"
            pub fn say() {{
                println!({});
            }}
    "#,
        args
    )
        .parse()
        .unwrap()
}

// bang 宏
#[proc_macro]
pub fn hashmap(input: TokenStream) -> TokenStream {
    let input = input.to_string();
    let input = input.trim_end_matches(',');
    let input = input.split(',').map(|v| {
        let mut data = if v.contains(':') {
            v.split(":")
        } else {
            v.split("=>")
        };
        let (key, value) = (data.next().unwrap(), data.next().unwrap());
        format!("hm.insert({}, {})", key, value)
    }).collect::<Vec<_>>();
    let count = input.len();
    let tokens = format!(r#" {{
            let mut hm = ::std::collections::HashMap::with_capacity({});
            {}
            hm
        }}
    "#, count, input.iter().map(|v| format!("{};", v)).collect::<String>());
    tokens.parse().unwrap()
}

/// 编写New
#[proc_macro_derive(New)]
pub fn derive_new(input: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(input as DeriveInput);
    let result = match ast.data {
        Data::Struct(ref s) => {
            new_for_struct(&ast, &s.fields)
        }
        _ => panic!("doesn't work with unions yet!")
    };
    result.into()
}

fn new_for_struct(ast: &DeriveInput, fields: &Fields) -> proc_macro2::TokenStream {
    match fields {
        Fields::Named(fields) => {
            new_impl(ast, Some(&fields.named), true)
        }
        Fields::Unnamed(fields) => {
            new_impl(ast, Some(&fields.unnamed), true)
        }
        Fields::Unit => {
            new_impl(ast, None, false)
        }
    }
}

fn new_impl(ast: &DeriveInput,
            fields: Option<&syn::punctuated::Punctuated<syn::Field, syn::Token![,]>>,
            named: bool) -> proc_macro2::TokenStream {
    let struct_name = &ast.ident;
    let unit = fields.is_none();
    let empty = Default::default();
    let fields : Vec<_> = fields.unwrap_or(&empty)
        .iter().enumerate()
        .map(|(i, f)| FieldExt::new(f, i, named)).collect();
    let args = fields.iter().map(|f| f.as_arg());
    let inits = fields.iter().map(|f| f.as_init());
    let inits = if unit {
        quote!()
    } else if named {
        quote! { {#(#inits),*} }
    } else {
        quote! { {#(#inits),*} }
    };
    let (impl_generics, ty_generics, where_clause) =
        ast.generics.split_for_impl();
    let (new, doc) = (
        syn::Ident::new("new", proc_macro2::Span::call_site()),
        format!("Constructs a new `{}`", struct_name));
    quote! {
        impl #impl_generics #struct_name #ty_generics #where_clause {
            #[doc= #doc]
            pub fn #new(#(#args),*) -> Self {
                #struct_name #inits
            }
        }
    }
}

struct FieldExt<'a> {
    ty: &'a syn::Type,
    ident: syn::Ident,
    named: bool,
}

impl<'a> FieldExt<'a> {
    pub fn new(field: &'a syn::Field, idx: usize, named: bool) -> Self {
        FieldExt {
            ty: &field.ty,
            ident: if named {
                field.ident.clone().unwrap()
            } else {
                syn::Ident::new(&format!("f{}", idx),
                                proc_macro2::Span::call_site())
            },
            named,
        }
    }

    pub fn as_arg(&self) -> proc_macro2::TokenStream {
        let f_name = &self.ident;
        let ty = &self.ty;
        quote!(#f_name : #ty)
    }

    pub fn as_init(&self) -> proc_macro2::TokenStream {
        let f_name = &self.ident;
        let init = quote!(#f_name);
        if self.named {
            quote!(#f_name : #init)
        } else {
            quote!(#init)
        }
    }
}